from contextlib import asynccontextmanager
from typing import AsyncGenerator

# from sqlmodel.ext.asyncio.session import AsyncSession
from sqlalchemy.ext.asyncio import AsyncEngine, AsyncSession, create_async_engine
from sqlalchemy.orm import sessionmaker

# 创建异步数据库引擎（支持连接池优化和自动重试）
async_engine: AsyncEngine = create_async_engine(
    "postgresql+asyncpg://postgres:2kcx674N4hMYrY45@10.200.200.2:5432/gce_db",
    # 连接池配置（优化高并发场景）
    pool_size=20,  # 常驻连接数
    max_overflow=10,  # 临时连接数上限
    pool_timeout=30,  # 获取连接超时（秒）
    pool_recycle=1800,  # 连接最大存活时间
    pool_pre_ping=True,  # 连接健康检查（防止数据库连接超时断开）
    # 性能优化配置
    echo=False,  # SQL 语句日志输出
    future=True,  # SQLAlchemy 2.0+ 风格
    # 连接参数优化
    connect_args=({"command_timeout": 60, "server_settings": {"application_name": "DEMO_PROJECT"}}),  # 命令执行超时
)


# 创建异步会话工厂（绑定引擎，关闭自动行为）
AsyncSessionLocal = sessionmaker(
    bind=async_engine,
    class_=AsyncSession,
    autocommit=False,
    autoflush=False,
    expire_on_commit=False,  # 提交后不使对象过期，适合 API 场景
)


@asynccontextmanager
async def async_db_session() -> AsyncGenerator[AsyncSession, None]:
    """
    异步数据库会话上下文管理器

    自动处理事务提交/回滚，确保资源正确释放。
    在发生异常时自动回滚事务，在正常结束时提交事务。

    使用示例：
    ```python
    # 异常情况下的使用
    try:
        async with async_db_session() as session:
            user = User(name="test")
            session.add(user)
            raise Exception("模拟异常")
            # 自动回滚事务，不会保存到数据库
    except Exception as e:
        print(f"操作失败: {e}")
    ```

    Yields:
        AsyncSession: 异步数据库会话实例

    Raises:
        Exception: 数据库连接或操作失败
    """
    session = AsyncSessionLocal()
    try:
        yield session
        await session.commit()
    except Exception:
        await session.rollback()
        raise
    finally:
        await session.close()
